//
// Created by yk120 on 2024/3/1.
//

#ifndef NANOFAISS_CLUSTERING_H
#define NANOFAISS_CLUSTERING_H

#include <nanofaiss/Index.h>


#include <vector>

namespace faiss {

/** Class for clustering parameters.
 */
struct ClusteringParameters {
    /// number of iteration
    int niter = 25;

    /// redo cluserting
    int nredo = 1;


    /// fewer than this number will warning, fewer than 1 will exception
    int min_points_per_centroid = 39;

    /// to limit size of dataset
    int max_points_per_centroid = 256;

    int seed = 1234;

    // when the training set is encoded, batch size of the codec decoder
    size_t decode_block_size = 32768;

};


struct Clustering : ClusteringParameters {
    size_t d; /// size of the vector dimension
    size_t k; /// number of the centroid

    /// centroids (k * d)
    std::vector<float> centroids;

    Clustering(int d, int k);

    Clustering(int d, int k, const ClusteringParameters& cp);

    /** run k-means training
     * @param n         number of vector
     * @param x         training vectors, size n * d
     * @param index     index used for assignment
     * @param x_weight
     */
    virtual void train(
            idx_t n,
            const float *x,
            Index &index);

    /** run with encoded vectors
     * take a codec to decode the input vector
     * @param codec     codec used to decode the vectors (nullptr = vector are in fact floats)
     */
    void train_encoded(
            idx_t nx,
            const uint8_t *xin,
            const Index *codec,
            Index &index);

    /// post-process centroids after each update
    /// include L2 normalization and nearest integer rounding
    void post_process_centroids();

};


}

#endif // NANOFAISS_CLUSTERING_H
